import pandas as pd
import os
from finrl.config import INDICATORS
from stockstats import StockDataFrame as Sdf


def read_stock_data(file_path):
    """
    读取股票数据文件
    
    参数:
    file_path : str
        股票数据文件路径 (CSV格式)
        
    返回:
    pd.DataFrame
        股票数据
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"股票数据文件 {file_path} 不存在")
    
    # 读取CSV文件
    df = pd.read_csv(file_path)
    
    # 确保必要的列存在
    required_columns = ['date', 'open', 'high', 'low', 'close', 'volume', 'tic']
    for col in required_columns:
        if col not in df.columns:
            raise ValueError(f"缺少必要的列: {col}")
    
    # 转换日期格式
    df['date'] = pd.to_datetime(df['date'])
    
    return df


def add_technical_indicators(df, tech_indicator_list=INDICATORS):
    """
    为股票数据添加技术指标
    
    参数:
    df : pd.DataFrame
        股票数据
    tech_indicator_list : list
        技术指标列表
        
    返回:
    pd.DataFrame
        添加技术指标后的数据
    """
    df = df.copy()
    df = df.sort_values(by=['tic', 'date'])
    stock = Sdf.retype(df.copy())
    unique_ticker = stock.tic.unique()

    for indicator in tech_indicator_list:
        indicator_df = pd.DataFrame()
        for i in range(len(unique_ticker)):
            try:
                temp_indicator = stock[stock.tic == unique_ticker[i]][indicator]
                temp_indicator = pd.DataFrame(temp_indicator)
                temp_indicator['tic'] = unique_ticker[i]
                temp_indicator['date'] = df[df.tic == unique_ticker[i]]['date'].to_list()
                indicator_df = pd.concat([indicator_df, temp_indicator], ignore_index=True)
            except Exception as e:
                print(f"计算指标 {indicator} 时出错: {e}")
        
        df = df.merge(
            indicator_df[['tic', 'date', indicator]], 
            on=['tic', 'date'], 
            how='left'
        )
    
    df = df.sort_values(by=['date', 'tic'])
    return df


def save_stock_data(df, output_path):
    """
    保存处理后的股票数据
    
    参数:
    df : pd.DataFrame
        处理后的股票数据
    output_path : str
        输出文件路径
    """
    # 确保输出目录存在
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 保存为CSV文件
    df.to_csv(output_path, index=False)
    print(f"数据已保存至 {output_path}")


def process_single_stock(input_file, output_file, indicators=INDICATORS):
    """
    处理单个股票文件
    
    参数:
    input_file : str
        输入文件路径
    output_file : str
        输出文件路径
    indicators : list
        技术指标列表
    """
    print(f"正在处理股票文件: {input_file}")
    
    # 读取数据
    df = read_stock_data(input_file)
    print(f"原始数据形状: {df.shape}")
    
    # 添加技术指标
    df_with_indicators = add_technical_indicators(df, indicators)
    print(f"添加指标后数据形状: {df_with_indicators.shape}")
    
    # 保存数据
    save_stock_data(df_with_indicators, output_file)
    
    print(f"完成处理: {input_file} -> {output_file}")


def process_multiple_stocks(input_dir, output_dir, indicators=INDICATORS):
    """
    批量处理多个股票文件
    
    参数:
    input_dir : str
        输入目录路径
    output_dir : str
        输出目录路径
    indicators : list
        技术指标列表
    """
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 获取所有CSV文件
    files = [f for f in os.listdir(input_dir) if f.endswith('.csv')]
    
    if not files:
        print(f"目录 {input_dir} 中没有找到CSV文件")
        return
    
    print(f"找到 {len(files)} 个股票文件待处理")
    
    for i, file in enumerate(files):
        print(f"处理进度: {i+1}/{len(files)}")
        input_file = os.path.join(input_dir, file)
        output_file = os.path.join(output_dir, file)
        
        try:
            process_single_stock(input_file, output_file, indicators)
        except Exception as e:
            print(f"处理文件 {file} 时出错: {e}")


# 示例用法
if __name__ == "__main__":
    # 示例1: 处理单个股票文件
    # process_single_stock(
    #     input_file="data/stock_sample.csv", 
    #     output_file="data/processed/stock_sample_with_indicators.csv"
    # )
    
    # 示例2: 批量处理多个股票文件
    # process_multiple_stocks(
    #     input_dir="data/stocks", 
    #     output_dir="data/stocks_with_indicators"
    # )
    
    # 创建示例数据用于测试
    sample_data = pd.DataFrame({
        'date': pd.date_range('2023-01-01', periods=100, freq='D'),
        'open': [100 + i*0.1 for i in range(100)],
        'high': [102 + i*0.1 for i in range(100)],
        'low': [98 + i*0.1 for i in range(100)],
        'close': [100 + i*0.1 for i in range(100)],
        'volume': [1000 + i*10 for i in range(100)],
        'tic': ['SAMPLE'] * 100
    })
    
    # 创建示例数据目录
    if not os.path.exists("example_data"):
        os.makedirs("example_data")
    
    # 保存示例数据
    sample_data.to_csv("example_data/sample_stock.csv", index=False)
    print("已创建示例数据文件: example_data/sample_stock.csv")
    
    # 处理示例数据
    process_single_stock(
        input_file="example_data/sample_stock.csv",
        output_file="example_data/sample_stock_with_indicators.csv"
    )